from stl import mesh
import numpy as np
import os
import tqdm
import cv2

BYTES_PER_PIXEL = 2
LINES = 21481
LINE_SAMPLES = 27288
scale_factor = 20
base_batch = 1000
rate = 1


def create_plane_segment(dem_image, row, col, output_path, rate=2):
    # 生成球面片段的顶点
    assert dem_image.shape[0] == dem_image.shape[1]
    vertices = []
    num_steps = dem_image.shape[0] // rate
    X = np.arange(
        row * dem_image.shape[0] * scale_factor,
        (row + 1) * dem_image.shape[0] * scale_factor,
        rate * scale_factor,
    )
    Y = np.arange(
        col * dem_image.shape[1] * scale_factor,
        (col + 1) * dem_image.shape[1] * scale_factor,
        rate * scale_factor,
    )

    dem_image = dem_image[::rate, ::rate]
    for i in range(num_steps):
        for j in range(num_steps):
            x = X[i] - row * rate * scale_factor
            y = Y[j] - col * rate * scale_factor
            z = dem_image[i, j]
            vertices.append((x, y, z))

    # 生成三角形
    faces = []
    for i in range(num_steps - 1):
        for j in range(num_steps - 1):
            # 当前点的索引
            current = i * num_steps + j
            # 下一个theta值的点的索引
            next_theta = current + num_steps
            faces.append([current, next_theta, current + 1])
            faces.append([current + 1, next_theta, next_theta + 1])

    # 创建网格
    sphere_segment = mesh.Mesh(np.zeros(len(faces), dtype=mesh.Mesh.dtype))
    for i, f in enumerate(faces):
        for j in range(3):
            sphere_segment.vectors[i][j] = vertices[f[j]]

        # 保存为STL文件
    sphere_segment.save(output_path)


if __name__ == "__main__":
    path = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/CE2_GRAS_DEM_20m_H003_07S135W_A.tif"
    output_dir = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"
    num_batches = LINES // base_batch
    num_samples = LINE_SAMPLES // base_batch
    # 参数
    print(
        f"The basic batch of DEM is {base_batch}x{base_batch}. Each line contains {num_samples} batches. There are totally {num_batches*num_samples} basic batches."
    )
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    for row in tqdm.tqdm(range(0, LINES // base_batch)):
        if not os.path.exists(os.path.join(output_dir, "stl_models")):
            os.makedirs(os.path.join(output_dir, "stl_models"))
        if not os.path.exists(os.path.join(output_dir, "stl_models", f"{row}")):
            os.makedirs(os.path.join(output_dir, "stl_models", f"{row}"))
        for col in range(0, LINE_SAMPLES // base_batch):
            img_array = img[
                row * base_batch : (row + 1) * base_batch + rate,
                col * base_batch : (col + 1) * base_batch + rate,
            ]
            num_steps = base_batch // rate
            create_plane_segment(
                img_array,
                row,
                col,
                os.path.join(
                    output_dir,
                    "stl_models",
                    f"{row}",
                    f"{row}_{col}_{num_steps}.stl",
                ),
                rate,
            )
            np.savez(
                os.path.join(
                    output_dir,
                    "dem_int16",
                    f"{row}",
                    f"dem_{row}_{col}_{num_steps}.npz",
                ),
                image=img_array[:-rate, :-rate],
            )
